import cv2
import numpy as np
from lane_detection.CameraTransform import CameraTransform
from lane_detection.read_param import CameraTransformation

class CameraCombine:
    """相机组合处理类。
    
    该类负责处理左右两个相机的图像，计算鸟瞰图(BEV)变换矩阵，并将两个相机的BEV图像融合。
    
    Attributes:
        left_camera (CameraTransform): 左相机变换对象。
        right_camera (CameraTransform): 右相机变换对象。
        bev_camera_center_mm (tuple): BEV相机中心位置(毫米)。
        mid_x (int): BEV图像中点X坐标。
    """
    
    def __init__(self, 
                 left_param: CameraTransformation,
                 right_param: CameraTransformation,
                 cal_outer_param: bool = True,
                 ang_bias: float = 5.0):
        """初始化CameraCombine对象。
        
        计算两个相机的BEV变换矩阵，包括相机位置、旋转角度和透视变换。
        
        Args:
            left_param (CameraTransformation): 左相机变换参数。
            right_param (CameraTransformation): 右相机变换参数。
            cal_outer_param (bool): 是否计算外参映射表。
            ang_bias (float): 图像偏航角偏差(度)。
        """
        # 初始化两个相机的CameraTransform对象
        self.left_camera = CameraTransform(left_param)
        self.right_camera = CameraTransform(right_param)

        # 计算两个相机在棋盘格坐标系中的位置
        R_inv_L = self.left_camera.rotation_matrix.T
        camera_in_world_L = -R_inv_L @ self.left_camera.position

        R_inv_R = self.right_camera.rotation_matrix.T
        camera_in_world_R = -R_inv_R @ self.right_camera.position

        # 计算BEV中心点（两个相机连线的中点）
        bev_center_x = (camera_in_world_L[0, 0] + camera_in_world_R[0, 0]) / 2
        bev_center_y = (camera_in_world_L[1, 0] + camera_in_world_R[1, 0]) / 2 
        bev_center = np.array([bev_center_x, bev_center_y])

        # 计算两个相机连线的方向向量
        dx = camera_in_world_R[0, 0] - camera_in_world_L[0, 0]
        dy = camera_in_world_R[1, 0] - camera_in_world_L[1, 0]

        # 计算旋转角度（使连线平行于x轴）
        rotation_angle = -np.arctan2(dy, dx)  # 弧度

        # 创建旋转矩阵（绕z轴旋转）
        R_z = np.array([
            [np.cos(rotation_angle), -np.sin(rotation_angle), 0],
            [np.sin(rotation_angle), np.cos(rotation_angle), 0],
            [0, 0, 1]
        ])

        # 尝试加载预先计算好的映射表
        loaded_left = False
        loaded_right = False

        if not cal_outer_param:
            # 加载预先计算好的映射表
            loaded_left = self.left_camera.load_bev_mapping_tables(left_param.mapping_outer)
            loaded_right = self.right_camera.load_bev_mapping_tables(right_param.mapping_outer)
            
        if cal_outer_param or loaded_left is False or loaded_right is False:
            # 计算两个相机的BEV变换矩阵
            self.left_camera.calculate_bev_camera_pose(bev_center, global_rotation=R_z)
            self.right_camera.calculate_bev_camera_pose(bev_center, global_rotation=R_z)

            # 计算两个相机的BEV中心位置（毫米）
            self.bev_camera_center_mm = (self.left_camera.bev_pos_mm[0] + (self.right_camera.bev_pos_mm[0] - self.left_camera.bev_pos_mm[0]) * 0.5,
                                        self.left_camera.bev_pos_mm[1] + (self.right_camera.bev_pos_mm[1] - self.left_camera.bev_pos_mm[1]) * 0.5)

            # 计算两个相机的BEV透视变换矩阵
            self.left_camera.calculate_perspective_transform(bev_center, global_rotation=R_z, bev_camera_center=self.bev_camera_center_mm, image_yaw_bias=ang_bias)
            self.right_camera.calculate_perspective_transform(bev_center, global_rotation=R_z, bev_camera_center=self.bev_camera_center_mm, image_yaw_bias=ang_bias)

            # 预计算映射表
            self.left_camera.precompute_bev_maps()
            self.right_camera.precompute_bev_maps()

            # 保存映射表
            self.left_camera.save_bev_mapping_tables(left_param.mapping_outer)
            self.right_camera.save_bev_mapping_tables(right_param.mapping_outer)

            
        # 计算BEV图像的中点位置（用于裁剪）
        self.mid_x = int((self.left_camera.bev_pos[0] + self.right_camera.bev_pos[0]) / 2)

        print("初始化left和right相机的BEV变换完成")

    def create_cropped_bev(self, bev_img_left: np.ndarray, bev_img_right: np.ndarray) -> np.ndarray:
        """创建裁剪后的BEV图像。
        
        左侧相机只保留左边画幅，右侧相机只保留右边画幅，然后合并。
        
        Args:
            bev_img_left (np.ndarray): 左相机BEV图像。
            bev_img_right (np.ndarray): 右相机BEV图像。
            
        Returns:
            np.ndarray: 合并后的BEV图像。
        """
        # 创建左侧相机的裁剪BEV（只保留左侧）
        cropped_left = bev_img_left.copy()
        cropped_left[:, self.mid_x:] = 0  # 右侧部分置为黑色
        
        # 创建右侧相机的裁剪BEV（只保留右侧）
        cropped_right = bev_img_right.copy()
        cropped_right[:, :self.mid_x] = 0  # 左侧部分置为黑色
        
        # 合并两个裁剪后的图像
        combined_bev = cv2.add(cropped_left, cropped_right)
        
        return combined_bev
    
    def process_cameras(self, left_img: np.ndarray, right_img: np.ndarray) -> np.ndarray:
        """处理两个相机的图像并返回融合后的BEV图像。
        
        Args:
            left_img (np.ndarray): 左相机原始图像。
            right_img (np.ndarray): 右相机原始图像。
            
        Returns:
            np.ndarray: 融合后的BEV图像。
        """
        # 处理左相机图像
        left_img = self.left_camera.high_precision_undistort(left_img)
        left_bev = self.left_camera.project_points_with_kdtree(left_img)

        # 处理右相机图像
        right_img = self.right_camera.high_precision_undistort(right_img)
        right_bev = self.right_camera.project_points_with_kdtree(right_img)
        
        # 添加相机位置标记
        font = cv2.FONT_HERSHEY_SIMPLEX

        cv2.circle(left_bev, (int(self.left_camera.bev_pos[0]), int(self.left_camera.bev_pos[1])), 20, (0, 255, 0), -1)  # 左侧相机位置
        cv2.circle(right_bev, (int(self.right_camera.bev_pos[0]), int(self.right_camera.bev_pos[1])), 20, (0, 0, 255), -1)  # 右侧相机位置
        cv2.putText(left_bev, f"Left Camera", (int(self.left_camera.bev_pos[0]) - 300, int(self.left_camera.bev_pos[1]) + 30), font, 1, (0, 255, 0), 2)
        cv2.putText(right_bev, f"Right Camera", (int(self.right_camera.bev_pos[0]) + 300, int(self.right_camera.bev_pos[1]) + 30), font, 1, (0, 0, 255), 2)

        # 简单融合两张BEV图像
        combined_bev = self.create_cropped_bev(left_bev, right_bev)
        return combined_bev